#include "gwbase/encrypt/aes_helper.h"

#include <openssl/aes.h>
#include <cassert>
#include <zlib.h>
#include <gwbase/base/file_ops.h>
#include <gwbase/base/logging.h>
#include "gwbase/encrypt/digest.h"

namespace dicom {

static inline size_t GetAesMaxOutlen(size_t len) {
  return (len / AES_BLOCK_SIZE + 1) * AES_BLOCK_SIZE;
  ;
}

Aes256Encrypt::Aes256Encrypt(const ByteArray &key, size_t block_size)
    : ctx_(NULL), block_size_(block_size), outlen_(0) {
  assert(key.size() == 32);
  ctx_ = new EVP_CIPHER_CTX();
  EVP_CIPHER_CTX_init(ctx_);
  int ret = EVP_EncryptInit(ctx_, EVP_aes_256_ecb(), &key[0], NULL);
  assert(ret = 1);

  out_ = new uint8_t[GetAesMaxOutlen(block_size)];
}

Aes256Encrypt::~Aes256Encrypt() {
  if (ctx_ != NULL) {
    EVP_CIPHER_CTX_cleanup(ctx_);
    delete ctx_;
    ctx_ = NULL;

    delete[] out_;
    out_ = NULL;
  }
}

bool Aes256Encrypt::Update(const uint8_t *in, size_t inlen) {
  assert(inlen <= block_size_);
  int ret =
      EVP_EncryptUpdate(ctx_, out_, &outlen_, in, static_cast<int>(inlen));
  assert(ret == 1);
  return true;
}

bool Aes256Encrypt::Finalize() {
  int ret = EVP_EncryptFinal_ex(ctx_, out_, &outlen_);
  assert(ret == 1);
  return true;
}

Aes256Decrypt::Aes256Decrypt(const ByteArray &key, size_t block_size)
    : ctx_(NULL), block_size_(block_size), outlen_(0) {
  assert(key.size() == 32);
  ctx_ = new EVP_CIPHER_CTX();
  EVP_CIPHER_CTX_init(ctx_);
  int ret = EVP_DecryptInit(ctx_, EVP_aes_256_ecb(), &key[0], NULL);
  assert(ret = 1);

  out_ = new uint8_t[GetAesMaxOutlen(block_size)];
}

Aes256Decrypt::~Aes256Decrypt() {
  if (ctx_ != NULL) {
    EVP_CIPHER_CTX_cleanup(ctx_);
    delete ctx_;
    ctx_ = NULL;

    delete[] out_;
    out_ = NULL;
  }
}

bool Aes256Decrypt::Update(const uint8_t *in, size_t inlen) {
  int ret =
      EVP_DecryptUpdate(ctx_, out_, &outlen_, in, static_cast<int>(inlen));
  assert(ret == 1);
  return true;
}

bool Aes256Decrypt::Finalize() {
  int ret = EVP_DecryptFinal_ex(ctx_, out_, &outlen_);
  // assert (ret == 1);
  return (ret == 1);
}
Aes128Encrypt::Aes128Encrypt(const ByteArray &key, size_t block_size)
    : ctx_(NULL), block_size_(block_size), outlen_(0) {
  assert(key.size() == 16);
  ctx_ = new EVP_CIPHER_CTX();
  EVP_CIPHER_CTX_init(ctx_);
  int ret = EVP_EncryptInit(ctx_, EVP_aes_128_ecb(), &key[0], NULL);
  assert(ret = 1);

  out_ = new uint8_t[GetAesMaxOutlen(block_size)];
}

Aes128Encrypt::~Aes128Encrypt() {
  if (ctx_ != NULL) {
    EVP_CIPHER_CTX_cleanup(ctx_);
    delete ctx_;
    ctx_ = NULL;

    delete[] out_;
    out_ = NULL;
  }
}

bool Aes128Encrypt::Update(const uint8_t *in, size_t inlen) {
  assert(inlen <= block_size_);
  int ret =
      EVP_EncryptUpdate(ctx_, out_, &outlen_, in, static_cast<int>(inlen));
  assert(ret == 1);
  return true;
}

bool Aes128Encrypt::Finalize() {
  int ret = EVP_EncryptFinal_ex(ctx_, out_, &outlen_);
  assert(ret == 1);
  return true;
}

Aes128Decrypt::Aes128Decrypt(const ByteArray &key, size_t block_size)
    : ctx_(NULL), block_size_(block_size), outlen_(0) {
  assert(key.size() == 16);
  ctx_ = new EVP_CIPHER_CTX();
  EVP_CIPHER_CTX_init(ctx_);
  int ret = EVP_DecryptInit(ctx_, EVP_aes_128_ecb(), &key[0], NULL);
  assert(ret = 1);

  out_ = new uint8_t[GetAesMaxOutlen(block_size)];
}

Aes128Decrypt::~Aes128Decrypt() {
  if (ctx_ != NULL) {
    EVP_CIPHER_CTX_cleanup(ctx_);
    delete ctx_;
    ctx_ = NULL;

    delete[] out_;
    out_ = NULL;
  }
}

bool Aes128Decrypt::Update(const uint8_t *in, size_t inlen) {
  int ret =
      EVP_DecryptUpdate(ctx_, out_, &outlen_, in, static_cast<int>(inlen));
  assert(ret == 1);
  return true;
}

bool Aes128Decrypt::Finalize() {
  int ret = EVP_DecryptFinal_ex(ctx_, out_, &outlen_);
  // assert (ret == 1);
  return (ret == 1);
}

ByteArray AESHelper::Encrypt128(const ByteArray &key, const ByteArray &plain) {
  assert(plain.size() > 0);
  ByteArray cipher;

  Aes128Encrypt enc(key, plain.size());
  enc.Update(&plain[0], plain.size());
  cipher.insert(cipher.end(), enc.data(), enc.end());

  enc.Finalize();
  cipher.insert(cipher.end(), enc.data(), enc.end());

  return cipher;
}

ByteArray AESHelper::Decrypt128(const ByteArray &key, const ByteArray &cipher) {
  assert(cipher.size() > 0);
  ByteArray plain;

  Aes128Decrypt dec(key, cipher.size());
  dec.Update(&cipher[0], cipher.size());
  plain.insert(plain.end(), dec.data(), dec.end());

  if (dec.Finalize() == false) {
    plain.clear();
    return plain;
  }

  plain.insert(plain.end(), dec.data(), dec.end());

  return plain;
}

ByteArray AESHelper::Encrypt256(const ByteArray &key, const ByteArray &plain) {
  assert(plain.size() > 0);
  ByteArray cipher;

  Aes256Encrypt enc(key, plain.size());
  enc.Update(&plain[0], plain.size());
  cipher.insert(cipher.end(), enc.data(), enc.end());

  enc.Finalize();
  cipher.insert(cipher.end(), enc.data(), enc.end());

  return cipher;
}

ByteArray AESHelper::Decrypt256(const ByteArray &key, const ByteArray &cipher) {
  assert(cipher.size() > 0);
  ByteArray plain;

  Aes256Decrypt dec(key, cipher.size());
  dec.Update(&cipher[0], cipher.size());
  plain.insert(plain.end(), dec.data(), dec.end());

  if (dec.Finalize() == false) {
    plain.clear();
    return plain;
  }

  plain.insert(plain.end(), dec.data(), dec.end());

  return plain;
}

Status AESHelper::EncFile(const ByteArray &key, const std::string &in,
                          const std::string &out, IStoppingFlag *stop) {
  const int kBlockSize = 4096;
  SequentialFilePtr inf;
  WritableFilePtr outf;
  Status s;
  s = SequentialFile::Open(in, &inf);
  if (!s.ok()) return s;

  s = WritableFile::Open(out, true, &outf);
  if (!s.ok()) return s;

  uint8_t buf[kBlockSize];
  Aes256Encrypt enc(key, kBlockSize);

  size_t read_bytes;
  while (1) {
    if (stop != NULL && stop->IsStopping()) goto failed;

    if (!inf->Read(buf, sizeof(buf), &read_bytes).ok()) goto failed;

    if (read_bytes == 0) {
      enc.Finalize();
      if (!outf->Append(enc.data(), enc.len()).ok()) goto failed;

      return Status();
    }

    // n > 0
    enc.Update(buf, read_bytes);
    if (!outf->Append(enc.data(), enc.len()).ok()) goto failed;
  }

failed:
  return Status(Status::kOther, "write error");
}

static std::string BinaryHashToHexString(const unsigned char *hash, int len) {
  std::string out;

  // reverse enough space to gain a little speedup
  out.reserve(len * 2 + 1);
  for (int i = 0; i < len; ++i)
    out += Fmt("%02x", static_cast<int>(hash[i])).c_str();
  return out;
}

Status AESHelper::CompressAndEncFile(
    const ByteArray &key, const std::string &in, const std::string &out,
    std::string *enc_sha1, std::string *src_md5, IStoppingFlag *stop) {
  Status s;
  std::string errmsg;
  const int kBlockSize = 4096;

  SequentialFilePtr inf;
  WritableFilePtr outf;
  s = SequentialFile::Open(in, &inf);
  if (!s.ok()) return s;

  s = WritableFile::Open(out, true, &outf);
  if (!s.ok()) return s;

  uint8_t buf[kBlockSize];
  Aes256Encrypt cipher(key, kBlockSize);

  int ret, flush;
  unsigned have;
  z_stream strm;
  unsigned char outbuf[kBlockSize];
  strm.zalloc = Z_NULL;
  strm.zfree = Z_NULL;
  strm.opaque = Z_NULL;

  SHA_CTX sha1ctx;
  unsigned char sha1hash[20] = {0};
  SHA1_Init(&sha1ctx);
  MD5_CTX md5ctx;
  unsigned char md5hash[16] = {0};
  MD5_Init(&md5ctx);

  ::deflateInit(&strm, Z_DEFAULT_COMPRESSION);  // will add zlib header

  size_t read_bytes;
  while (1) {
    if (stop != NULL && stop->IsStopping()) {
      s = Status(Status::kOther, "user cancelled");
      goto error;
    }

    if (!inf->Read(buf, sizeof(buf), &read_bytes).ok()) {
      goto error;
    }

    MD5_Update(&md5ctx, buf, read_bytes);

    flush = (read_bytes == 0 ? Z_FINISH : Z_NO_FLUSH);
    strm.avail_in = static_cast<uInt>(read_bytes);
    strm.next_in = buf;

    do {
      strm.avail_out = kBlockSize;
      strm.next_out = outbuf;
      ret = ::deflate(&strm, flush);
      assert(ret != Z_STREAM_ERROR);
      have = kBlockSize - strm.avail_out;
      switch (ret) {
        case Z_NEED_DICT:
        case Z_DATA_ERROR:
        case Z_MEM_ERROR:
          s = Status(Status::kOther, "zlib::inflate() failed");
          goto error;
      }

      if (!cipher.Update(outbuf, have)) {
        s = Status(Status::kOther, "AES256.Update() failed");
        goto error;
      }

      if (cipher.len() > 0) {
        if (!outf->Append(cipher.data(), cipher.len()).ok()) {
          // errmsg = ::StrError();
          // LOG_NORMAL("failed to write output: %s : %s", utf8(out).c_str(),
          // errmsg.c_str());
          s = Status(Status::kOther, errmsg);
          goto error;
        }

        SHA1_Update(&sha1ctx, cipher.data(), cipher.len());
      }
    } while (strm.avail_out == 0);

    if (ret == Z_STREAM_END) break;
  }

  if (!cipher.Finalize()) {
    s = Status(Status::kOther, "AES256.Finalize() failed");
    goto error;
  }

  if (cipher.len() > 0) {
    if (!outf->Append(cipher.data(), cipher.len()).ok()) {
      // errmsg = ::StrError();
      // LOG_NORMAL("failed to write output: %s : %s", utf8(out).c_str(),
      // errmsg.c_str());
      s = Status(Status::kOther, errmsg);
      goto error;
    }

    SHA1_Update(&sha1ctx, cipher.data(), cipher.len());
  }

  ::deflateEnd(&strm);
  SHA1_Final(sha1hash, &sha1ctx);
  enc_sha1->assign(BinaryHashToHexString(sha1hash, 20));
  MD5_Final(md5hash, &md5ctx);
  src_md5->assign(BinaryHashToHexString(md5hash, 16));
  return Status();

error:
  ::deflateEnd(&strm);
  return s;
}

Status AESHelper::DecAndUncompressFile(const ByteArray &key,
                                       const std::string &in,
                                       const std::string &out,
                                       IStoppingFlag *stop) {
  Status s;
  std::string errmsg;
  const int kBlockSize = 4096;
  SequentialFilePtr inf;
  WritableFilePtr outf;
  uint8_t buf[kBlockSize];
  Aes256Decrypt dec(key, kBlockSize);

  int ret, flush;
  unsigned have;
  z_stream strm;
  unsigned char outbuf[kBlockSize];
  strm.zalloc = Z_NULL;
  strm.zfree = Z_NULL;
  strm.opaque = Z_NULL;

  ::inflateInit(&strm);

  s = SequentialFile::Open(in, &inf);
  if (!s.ok()) return s;

  s = WritableFile::Open(out, true, &outf);
  if (!s.ok()) return s;

  size_t read_bytes;
  while (1) {
    if (stop != NULL && stop->IsStopping()) {
      s = Status(Status::kOther, "user cancelled enc and uncompress");
      goto error;
    }

    if (!inf->Read(buf, sizeof(buf), &read_bytes).ok()) {
      // errmsg = ::StrError();
      // LOG_NORMAL("fread() failed: %s : %s", utf8(out).c_str(),
      // errmsg.c_str());
      s = Status(Status::kOther, "fread() failed");
      goto error;
    }

    flush = 0;
    if (read_bytes == 0) {
      if (!dec.Finalize()) {
        s = Status(Status::kOther, "AES256.Finalize() failed");
        goto error;
      }
      flush = Z_FINISH;
    } else {
      if (!dec.Update(buf, read_bytes)) {
        s = Status(Status::kOther, "AES256.Update() failed");
        goto error;
      }
      flush = Z_NO_FLUSH;
    }
    dec.data();
    strm.avail_in = dec.len();
    strm.next_in = const_cast<Bytef *>(dec.data());

    do {
      strm.avail_out = kBlockSize;
      strm.next_out = outbuf;
      ret = inflate(&strm, flush);
      assert(ret != Z_STREAM_ERROR);
      have = kBlockSize - strm.avail_out;
      switch (ret) {
        case Z_NEED_DICT:
        case Z_DATA_ERROR:
        case Z_MEM_ERROR:
          s = Status(Status::kOther, "zlib::inflate() failed");
          goto error;
      }

      if (!outf->Append(outbuf, have).ok()) {
        // errmsg = ::StrError();
        // LOG_NORMAL("failed to write output: %s : %s", utf8(out).c_str(),
        // errmsg.c_str());
        s = Status(Status::kOther, errmsg);
        goto error;
      }
    } while (strm.avail_out == 0);

    if (ret == Z_STREAM_END) break;
  }

  ::inflateEnd(&strm);
  return Status();

error:
  ::inflateEnd(&strm);
  return s;
}

Status AESHelper::DecFile(const ByteArray &key, const std::string &in,
                          const std::string &out, IStoppingFlag *stop) {
  const int kBlockSize = 4096;

  SequentialFilePtr inf;
  WritableFilePtr outf;
  Status s;
  s = SequentialFile::Open(in, &inf);
  if (!s.ok()) return s;

  s = WritableFile::Open(out, true, &outf);
  if (!s.ok()) return s;

  uint8_t buf[kBlockSize];
  Aes256Decrypt dec(key, kBlockSize);

  size_t read_bytes;
  while (1) {
    if (stop != NULL && stop->IsStopping()) goto failed;

    if (!inf->Read(buf, sizeof(buf), &read_bytes).ok()) goto failed;

    if (read_bytes == 0) {
      dec.Finalize();
      if (!outf->Append(dec.data(), dec.len()).ok()) goto failed;

      return Status();
    }

    // n > 0
    dec.Update(buf, read_bytes);
    if (!outf->Append(dec.data(), dec.len()).ok()) goto failed;
  }

failed:
  return Status(Status::kOther, "error");
}

// 生成用于加密、解密private key的密钥
static ByteArray GenAesKeyFromPlainPassword(const std::string &plain_pwd) {
  ByteArray pk_enc_key = ByteArray::fromRawStr(plain_pwd);
  size_t to_pad = (plain_pwd.size() >= 32 ? 0 : 32 - plain_pwd.size());
  pk_enc_key.resize(32);
  for (size_t i = 0; i < to_pad; ++i) pk_enc_key[32 - i - 1] = 0;

  return pk_enc_key;
}

ByteArray AESHelper::DecryptPrivateKey(const std::string &plain_pwd,
                                       const ByteArray &pk_cipher) {
  ByteArray pk_enc_key = GenAesKeyFromPlainPassword(plain_pwd);
  return AESHelper::Decrypt256(pk_enc_key, pk_cipher);
}

ByteArray AESHelper::EncryptPrivateKey(const std::string &plain_pwd,
                                       const ByteArray &pk_plain) {
  ByteArray pk_enc_key = GenAesKeyFromPlainPassword(plain_pwd);
  return AESHelper::Encrypt256(pk_enc_key, pk_plain);
}
}
